import os
from pathlib import Path

# 原始数据集路径
existing_labels_dir = Path("D:\Datasets\dataset2CheckNew\labels")

# 新数据集路径
detector_labels_dir = Path("D:\Datasets\dataset2CheckNew\labels-helwheel")

# 创建新目录
detector_labels_dir.mkdir(exist_ok=True, parents=True)

# 类别映射：旧ID -> 新ID
# 假设原始:  hel=0, nohel=1, wheel=2
# 新的: helmet=0,wheel=1,
class_mapping = {
    1: 0,  # nohel -> helmet
    2: 1  # wheel -> wheel
    # hel (0) remains 0, no need to map
}

# 转换标签
for label_file in existing_labels_dir.glob("*.txt"):
    new_lines = []
    with open(label_file, 'r') as f:
        for line in f:
            parts = line.strip().split()
            if not parts:
                continue
            old_class_id = int(parts[0])
            # 只处理需要转换的类别
            if old_class_id in class_mapping:
                new_class_id = class_mapping[old_class_id]
                parts[0] = str(new_class_id)
                new_lines.append(" ".join(parts))
            elif old_class_id == 0:  # wheel, ID unchanged
                new_lines.append(line.strip())  # keep original line
            # 忽略其他类别（如果有）

    # 写入新文件
    new_label_path = detector_labels_dir / label_file.name
    with open(new_label_path, 'w') as f:
        f.write("\n".join(new_lines))

print("标签转换完成！")
